package com.km.config;

import brave.Tracing;
import brave.context.slf4j.MDCCurrentTraceContext;
import brave.http.HttpTracing;
import brave.sampler.Sampler;
import brave.spring.web.TracingClientHttpRequestInterceptor;
import brave.spring.webmvc.TracingHandlerInterceptor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Import;
import org.springframework.http.client.ClientHttpRequestInterceptor;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.servlet.config.annotation.InterceptorRegistry;
import org.springframework.web.servlet.config.annotation.WebMvcConfigurerAdapter;
import zipkin.Span;
import zipkin.reporter.AsyncReporter;
import zipkin.reporter.Reporter;
import zipkin.reporter.Sender;
import zipkin.reporter.okhttp3.OkHttpSender;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;

/**
 * <p></p>
 * Created by zhezhiyong@163.com on 2017/8/29.
 */
@Configuration
// Importing these classes is effectively the same as declaring bean methods
@Import({TracingClientHttpRequestInterceptor.class, TracingHandlerInterceptor.class})
public class TracingConfiguration extends WebMvcConfigurerAdapter {

    @Bean
    RestTemplate template() {
        return new RestTemplate();
    }

    /**
     * Controls aspects of tracing such as the name that shows up in the UI
     */
    @Bean
    Tracing tracing() {
        Sender sender = OkHttpSender.create("http://192.168.97.57:9411/api/v1/spans");
        Reporter<Span> reporter = AsyncReporter.builder(sender).build();
        return Tracing.newBuilder()
                .localServiceName("demo-brave-springboot")
                .traceId128Bit(true)
                .sampler(Sampler.ALWAYS_SAMPLE)
//                .sampler(Sampler.create((float) 0.5)) //采样率
                .currentTraceContext(MDCCurrentTraceContext.create()) // puts trace IDs into logs
                .reporter(reporter).build();
    }

    // decides how to name and tag spans. By default they are named the same as the http method.
    @Bean
    HttpTracing httpTracing() {
        return HttpTracing.create(tracing());
    }

    @Autowired
    private TracingHandlerInterceptor serverInterceptor;

    @Autowired
    private TracingClientHttpRequestInterceptor clientInterceptor;

    @Autowired
    private RestTemplate restTemplate;

    @PostConstruct
    public void init() {
        List<ClientHttpRequestInterceptor> interceptors = new ArrayList<>(restTemplate.getInterceptors());
        interceptors.add(clientInterceptor);
        restTemplate.setInterceptors(interceptors);
    }

    @Override
    public void addInterceptors(InterceptorRegistry registry) {
        registry.addInterceptor(serverInterceptor);
    }
}